import os
import h5py
import csv

import pandas as pd


# ============================
# 任务1: 从 CSD 文件读取数据并转换为 CSV
# ============================

def convert_csd_to_csv(csd_file, output_csv):
    """
    读取 CSD 文件并将数据转换为 CSV 文件。

    参数:
    - csd_file: str，CSD 文件的路径，包含 CMU-MOSEI 数据的标签信息。
    - output_csv: str，输出的 CSV 文件路径，转换后的数据将保存到此文件中。

    处理流程:
    1. 打开并读取 CSD 文件（HDF5 格式）。
    2. 获取其中的 'All Labels' 组，以及 'data' 组。
    3. 遍历 'data' 组中的每个子数据集，提取 features 和 intervals 信息。
    4. 将每一行数据写入到 CSV 文件中，格式为：name + feature_0...feature_6 + interval_0, interval_1。
    """
    with h5py.File(csd_file, 'r') as lan:
        print(f"Keys in the HDF5 file: {list(lan.keys())}")

        # 访问 'All Labels' 组并打印其结构
        opinion_segment_labels_group = lan['All Labels']
        print(f"'All Labels' 组的结构: {list(opinion_segment_labels_group.keys())}")

        # 访问 'data' 组
        data_group = opinion_segment_labels_group['data']
        print(f"\n'data' 组的结构: {list(data_group.keys())}")
        print(f"数据总数: {len(data_group)}")

        # 创建 CSV 文件并写入数据
        with open(output_csv, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)

            # 写入 CSV 文件的标题行
            writer.writerow(['name'] + [f'feature_{i}' for i in range(7)] + [f'interval_{i}' for i in range(2)])

            # 遍历 'data' 组中的每个子数据集，写入到 CSV
            for idx, key in enumerate(data_group.keys()):
                data = data_group[key]

                # 获取 'features' 和 'intervals'
                features = data['features'][()]
                intervals = data['intervals'][()]

                # 生成带有后缀的 name
                if list(data_group.keys()).count(key) > 1:
                    key = f"{key}_{idx}"

                # 将数据写入 CSV 文件
                for i in range(len(features)):
                    row = [key] + list(features[i]) + list(intervals[i])
                    writer.writerow(row)
    print("CSV 文件已生成！")


# ============================
# 任务2: 从 TXT 文件中匹配时间戳并合并数据
# ============================

def match_txt_with_csv(csv_file, txt_directory, output_csv):
    """
    根据 TXT 文件中的时间戳匹配并合并 CSV 文件中的数据。

    参数:
    - csv_file: str，输入的 CSV 文件路径，包含 'name', 'feature_0' 至 'feature_6', 'interval_0', 'interval_1'。
    - txt_directory: str，TXT 文件所在的目录路径。TXT 文件包含每个视频段的时间戳和标签信息。
    - output_csv: str，输出的 CSV 文件路径，合并后的数据将保存到此文件中。

    处理流程:
    1. 读取 CSV 文件，并提取每一行的数据。
    2. 读取指定目录中的所有 TXT 文件，将每个文件中的时间戳和对应的视频名、编号存储到字典中。
    3. 对于每一行 CSV 数据，查找是否存在匹配的时间戳。
    4. 如果找到匹配项，则将 TXT 文件中的视频名和编号与 CSV 数据合并，生成新的数据行。
    5. 将合并后的数据保存到输出的 CSV 文件。
    """
    # 读取 CSV 文件
    df = pd.read_csv(csv_file)

    # 用于保存新的数据
    new_data = []

    # 读取 TXT 文件并存储为字典，以便快速查找
    txt_data = {}
    for filename in os.listdir(txt_directory):
        if filename.endswith(".txt"):
            with open(os.path.join(txt_directory, filename), 'r', encoding='utf-8') as f:
                for line in f:
                    parts = line.strip().split("___")
                    if len(parts) >= 4:
                        try:
                            name_part, number, interval_0_txt, interval_1_txt, _ = parts[:5]
                            txt_data[(float(interval_0_txt), float(interval_1_txt))] = (name_part, number)
                        except ValueError as e:
                            print(f"Error unpacking line: {line} - {e}")
                    else:
                        print(f"Skipping malformed line: {line}")

    # 遍历 CSV 中的每一行并匹配 TXT 数据
    for idx, row in df.iterrows():
        name = row['name']
        interval_0 = row['interval_0']
        interval_1 = row['interval_1']

        # 查找匹配的时间戳
        match = txt_data.get((interval_0, interval_1))

        if match:
            txt_name, number = match
            new_name = f"{txt_name}_{number}"

            # 提取 CSV 行中的 features
            features = row[
                ['feature_0', 'feature_1', 'feature_2', 'feature_3', 'feature_4', 'feature_5', 'feature_6']].tolist()

            # 新建数据行并加入新数据
            new_data.append([new_name] + features)

    # 创建新的 DataFrame 并保存为 CSV
    new_df = pd.DataFrame(new_data,
                          columns=['name', 'feature_0', 'feature_1', 'feature_2', 'feature_3', 'feature_4', 'feature_5',
                                   'feature_6'])
    new_df.to_csv(output_csv, index=False)
    print(f"新的 CSV 文件已生成: {output_csv}")


# ============================
# 任务3: 计算加权情感评分并映射为情感标签
# ============================


def calculate_emotion_scores(input_csv, output_csv):
    """
    计算加权情感评分，并将其映射到情感标签。

    参数:
    - input_csv: str，输入的 CSV 文件路径，包含特征数据和视频名称。
    - output_csv: str，输出的 CSV 文件路径，包含情感评分（0-5）的结果。
    """
    # 读取 CSV 文件
    df = pd.read_csv(input_csv)

    def cmumosei_round(a):
        if a < -2:
            return -3
        elif -2 <= a < -1:
            return -2
        elif -1 <= a < 0:
            return -1
        elif a == 0:  # 修正错误条件
            return 0
        elif 0 < a <= 1:
            return 1
        elif 1 < a <= 2:
            return 2
        else:
            return 3  # a > 2


    def cmumosei_round_small(a):
        if a <= -2.5:
            return -3
        elif -2.5 < a <= -1.5:
            return -2
        elif -1.5 < a <= -0.5:
            return -1
        elif -0.5 < a < 0.5:  # 修正错误条件
            return 0
        elif 0.5 <= a < 1.5:
            return 1
        elif 1.5 <= a < 2.5:
            return 2
        else:
            return 3  # a > 2

    def cmumosei_max(row):
        """获取 feature_1 到 feature_6 中的最大值所在的列号"""
        features = [f"feature_{i}" for i in range(1, 7)]
        max_key = max(features, key=lambda k: row[k])
        return int(max_key.split('_')[1])

    # df["emotion_score"] = df.apply(cmumosei_max, axis=1)

    # 生成情感得分列
    df["emotion_score"] = df["feature_0"].apply(cmumosei_round)
    # 将 -3~3 变换到 0~6
    df["emotion_score"] = df["emotion_score"] + 3

    # 10. 保存结果到 CSV 文件
    df.to_csv(output_csv, index=False)
    print(f"情感得分已生成并保存到 {output_csv}")



# ============================
# 任务4: 合并视频名和情感标签
# ============================

def merge_labels_with_text(input_text_csv, input_emotion_csv, output_csv):
    """
    将视频名称与情感标签合并，并保存最终的结果。

    参数:
    - input_text_csv: str，包含视频名称和英文文本的 CSV 文件路径。
    - input_emotion_csv: str，包含视频名称和情感评分的 CSV 文件路径。
    - output_csv: str，输出的 CSV 文件路径，包含合并后的视频名称与情感标签。

    处理流程:
    1. 读取文本数据文件（包含视频名称和文本内容）。
    2. 读取情感评分文件（包含视频名称和情感评分）。
    3. 根据视频名称合并两个数据集。
    4. 清理数据，去除包含缺失值的行。
    5. 删除特征列（'feature_0' 到 'feature_6'）。
    6. 保存最终合并的数据到新的 CSV 文件中。
    """
    # 读取原始文本数据
    df_text = pd.read_csv(input_text_csv)

    # 读取情感评分数据
    df_emotion = pd.read_csv(input_emotion_csv)

    # 合并数据，基于 'name' 字段
    df_merged = pd.merge(df_text, df_emotion, on='name', how='left')

    # 删除包含 NaN 值的行
    df_merged_clean = df_merged.dropna()

    # 显式地进行四舍五入并转换为整数类型
    df_merged_clean['emotion_score'] = df_merged_clean['emotion_score'].round().astype(int)

    # 将列名从 'emotion_score' 改为 'label'
    df_merged_clean.rename(columns={'emotion_score': 'label'}, inplace=True)

    # 删除 'feature_0' 到 'feature_6' 列，只保留视频名称和标签
    df_merged_clean.drop(columns=[f'feature_{i}' for i in range(7)], inplace=True)

    # 保存合并后的数据
    df_merged_clean.to_csv(output_csv, index=False)
    print(f"清理后的数据已保存到 {output_csv}")



# ============================
# 执行流程
# ============================

if __name__ == "__main__":
    # 执行任务1: 从 CSD 文件转换为 CSV
    print("开始执行任务1：从 CSD 文件转换为 CSV")
    convert_csd_to_csv('../../data/cmu_mosei/CMU_MOSEI_Labels.csd', '../../data/cmu_mosei/mosei_csd.csv')

    # 执行任务2: 从 TXT 文件匹配时间戳并合并数据
    print("开始执行任务2：从 TXT 文件匹配时间戳并合并数据")
    match_txt_with_csv('../../data/cmu_mosei/mosei_csd.csv', '../../data/Raw/Transcript/Segmented/Combined', '../../data/cmu_mosei/new_features.csv')

    # 执行任务3: 计算加权情感评分
    print("开始执行任务3：计算加权情感评分")
    calculate_emotion_scores('../../data/cmu_mosei/new_features.csv', '../../data/cmu_mosei/emotion_scores.csv')

    # 执行任务4: 合并视频名和情感标签
    print("开始执行任务4：合并视频名和情感标签")
    merge_labels_with_text('cmu_mosei.csv', '../../data/cmu_mosei/emotion_scores.csv', '../../data/cmu_mosei/merged_data_clean_max.csv')

    print("所有任务执行完毕！")
